import cv2
import numpy as np
from math import floor, ceil
from typing import List, Tuple


def show_img(title, frame):
    cv2.namedWindow(title, cv2.WINDOW_NORMAL)
    cv2.imshow(title, frame)
    cv2.waitKey(1)


def find_contours(frame: np.ndarray, gray=True, blur=((5, 5), 0), thresh=(255, 1, 1, 11, 2)):
    # show_img('frame', frame)
    gray = gray and len(frame.shape) == 3
    gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) if gray else frame
    blur_frame = cv2.GaussianBlur(gray_frame, *blur) if blur is not None else gray_frame
    thresh_frame = cv2.adaptiveThreshold(blur_frame, *thresh) if thresh is not None else blur_frame
    # show_img('thresh', thresh_frame)
    cnts, _ = cv2.findContours(thresh_frame, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
    return cnts


def extract_min_circle(frame: np.ndarray, **kwargs) -> List[Tuple[Tuple, float]]:
    """
    function: extract objects' min circle
    :param frame: image frame
    :return: a list of center and radius of extracted objects, like [((x, y), r), ]
    """
    return [cv2.minEnclosingCircle(cnt) for cnt in find_contours(frame, **kwargs)]


def extract_min_rect(frame: np.ndarray, **kwargs) -> List[Tuple[float, float, float, float]]:
    """
    function: extract objects' min rectangle
    :param frame: image frame
    :return: a list of center and radius of extracted objects, like [(x, y, w, h), ]
    """
    return [cv2.boundingRect(cnt) for cnt in find_contours(frame, **kwargs)]


def show_obj(frame: np.ndarray, center: Tuple, width: float, height: float, title: str):
    print('object %s: %s %f %f' % (title, center, width, height))
    v1_x, v1_y = floor(center[0] - width / 2), floor(center[1] - height / 2)
    v2_x, v2_y = ceil(center[0] + width / 2), ceil(center[1] + height / 2)
    roi = frame[max(v1_y, 0): max(v2_y, 1), max(v1_x, 0): max(v2_x, 1)]
    show_img(title, roi)


if __name__ == '__main__':
    img = cv2.imread("sample.png")
    extract_type = input('1: circle, 2: rectangle\n')
    if extract_type == '1':
        objects = extract_min_circle(img)
        print('%d object detected' % len(objects))
        for (index, (c, r)) in enumerate(objects):
            show_obj(img, c, r * 2, r * 2, str(index))
    else:
        objects = extract_min_rect(img)
        print('%d object detected' % len(objects))
        for (index, (x, y, w, h)) in enumerate(objects):
            show_obj(img, (x+w/2, y+h/2), w, h, str(index))
    if cv2.waitKey(0) == ord('q'):
        cv2.destroyAllWindows()
